package medium

import (
	"container/heap"
)

/*
时间复杂度:O(NlogN)
空间复杂度:O(logN)

问题:如何将学生安排到对应的班级,能使得总通过率最大??
优先队列:能便捷的对动态数据进行处理

将每个班级通过率a/b, 加上一个一定通过的学生后,带来的影响是x=(a+1/b+1)-a/b,而我们将学生安排到的那个班级,应该是x最大的那个班级
所以构建一个大根堆,根据x的值进行堆排序,x值大的在堆顶

糖水不等式: (a+k)/(b+k) > a/b  (a>0,b>0,k>0)
*/

func maxAverageRatio(classes [][]int, extraStudents int) float64 {
	h := &MyHeap{
		IsMaxHeap: true,
	}
	heap.Init(h)
	for i, v := range classes {
		bf := benefit(v[0], v[1])
		heap.Push(h, Pair{
			Index: i,
			Val:   bf,
		})
	}
	for extraStudents > 0 {
		p := heap.Pop(h).(Pair)
		classes[p.Index][0]++
		classes[p.Index][1]++
		p.Val = benefit(classes[p.Index][0], classes[p.Index][1])
		heap.Push(h, p)
		extraStudents--
	}
	var result float64
	for _, v := range classes {
		result += float64(v[0]) / float64(v[1])
	}
	return result / float64(len(classes))
}

func benefit(a, b int) float64 {
	return float64(a+1)/float64(b+1) - float64(a)/float64(b)
}

type Pair struct {
	Index int
	Val   float64
}

type MyHeap struct {
	Arr       []Pair
	IsMaxHeap bool
}

// 计算序列长度
func (h MyHeap) Len() int {
	return len(h.Arr)
}

// h.Arr[i] < h.Arr[j]表示小根堆，h.Arr[i] > h.Arr[j]表示大根堆
func (h MyHeap) Less(i, j int) bool {
	if !h.IsMaxHeap {
		// 使表达式为true的i所对应元素排前面，j所对应元素排后面；否则i所对应元素排后，j所对应元素排前
		return h.Arr[i].Val < h.Arr[j].Val
	} else {
		return h.Arr[i].Val > h.Arr[j].Val
	}
}

// 交换元素
func (h MyHeap) Swap(i, j int) {
	h.Arr[i], h.Arr[j] = h.Arr[j], h.Arr[i]
}

// 不要直接调用h.Push和Pop方法，而是调用heap.Push()和heap.Pop()
// heap.Push()和heap.Pop()内部会调用h.Push()和h.Pop()进行元素的添加和删除
// 然后再调用up(Push)或down(Pop)对重新调整顺序
func (h *MyHeap) Push(x interface{}) {
	(*h).Arr = append((*h).Arr, x.(Pair))
}

// 删除元素
func (h *MyHeap) Pop() interface{} {
	n := len(h.Arr) - 1
	x := h.Arr[n]
	h.Arr = h.Arr[:n]
	return x
}
